const fs = require('fs');
const { getUbyteFileData } = require('./parseMinist');

// 激活函数 sigmoid，用于引入非线性，将输入值映射到 (0, 1) 区间
function sigmoid(x) {
  return 1 / (1 + Math.exp(-x));
}

// 将一维数组转换为二维数组，方便后续卷积和池化操作
// 参数 data 是一维数组，width 是二维数组每一行的元素个数
function convert1Dto2D(data, width) {
  // 计算二维数组的行数
  const height = data.length / width;
  const result = [];
  // 遍历每一行
  for (let i = 0; i < height; i++) {
    // 从一维数组中截取对应行的数据添加到二维数组中
    result.push(data.slice(i * width, (i + 1) * width));
  }
  return result;
}

// 卷积操作，将输入数据与卷积核进行卷积运算
// 参数 input 是输入的二维数据，kernel 是卷积核，stride 是卷积步长
function convolve(input, kernel, stride = 1) {
  // 获取输入数据的高度和宽度
  const inputHeight = input.length;
  const inputWidth = input[0].length;
  // 获取卷积核的高度和宽度
  const kernelHeight = kernel.length;
  const kernelWidth = kernel[0].length;
  // 计算卷积输出的高度和宽度
  const outputHeight = Math.floor((inputHeight - kernelHeight) / stride) + 1;
  const outputWidth = Math.floor((inputWidth - kernelWidth) / stride) + 1;

  const output = [];
  // 遍历输出的每一行
  for (let y = 0; y < outputHeight; y++) {
    output[y] = [];
    // 遍历输出的每一列
    for (let x = 0; x < outputWidth; x++) {
      let sum = 0;
      // 遍历卷积核的每一行
      for (let ky = 0; ky < kernelHeight; ky++) {
        // 遍历卷积核的每一列
        for (let kx = 0; kx < kernelWidth; kx++) {
          // 计算卷积核元素与输入数据对应位置元素的乘积并累加
          sum += input[y * stride + ky][x * stride + kx] * kernel[ky][kx];
        }
      }
      // 将累加结果存入输出数组
      output[y][x] = sum;
    }
  }
  return output;
}

// 最大池化操作，对输入数据进行下采样，减少数据维度
// 参数 input 是输入的二维数据，poolSize 是池化窗口大小，stride 是池化步长
function maxPooling(input, poolSize = 2, stride = 2) {
  // 获取输入数据的高度和宽度
  const inputHeight = input.length;
  const inputWidth = input[0].length;
  // 计算池化输出的高度和宽度
  const outputHeight = Math.floor((inputHeight - poolSize) / stride) + 1;
  const outputWidth = Math.floor((inputWidth - poolSize) / stride) + 1;

  const output = [];
  // 遍历输出的每一行
  for (let y = 0; y < outputHeight; y++) {
    output[y] = [];
    // 遍历输出的每一列
    for (let x = 0; x < outputWidth; x++) {
      let maxValue = -Infinity;
      // 遍历池化窗口的每一行
      for (let py = 0; py < poolSize; py++) {
        // 遍历池化窗口的每一列
        for (let px = 0; px < poolSize; px++) {
          // 获取池化窗口内的元素值
          const value = input[y * stride + py][x * stride + px];
          // 如果该值大于当前最大值，则更新最大值
          if (value > maxValue) {
            maxValue = value;
          }
        }
      }
      // 将最大值存入输出数组
      output[y][x] = maxValue;
    }
  }
  return output;
}

// 扁平化操作，将二维数组转换为一维数组，以便输入到全连接层
function flatten(input) {
  return input.flat();
}

// 全连接层前向传播，将输入数据与权重相乘并加上偏置，然后通过 sigmoid 函数进行激活
// 参数 input 是输入的一维数据，weights 是全连接层的权重矩阵，biases 是偏置向量
function fullyConnected(input, weights, biases) {
  const output = [];
  // 遍历输出的每个元素
  for (let i = 0; i < weights.length; i++) {
    let sum = biases[i];
    // 遍历输入的每个元素
    for (let j = 0; j < input.length; j++) {
      // 计算输入元素与对应权重的乘积并累加到总和中
      sum += input[j] * weights[i][j];
    }
    // 对总和应用 sigmoid 激活函数并存入输出数组
    output[i] = sigmoid(sum);
  }
  return output;
}

// 从文件加载模型
function loadModel(filePath) {
    try {
        const modelData = fs.readFileSync(filePath, 'utf8');
        const model = JSON.parse(modelData);
        console.log(`Model loaded from ${filePath}`);
        return model;
    } catch (error) {
        console.error('Error loading model:', error);
        return null;
    }
}

function log(name, data) {
  fs.writeFileSync(name, JSON.stringify(data, null, 2));
}

// 预测函数，使用加载的模型对输入图片进行预测
function predict(image, model) {
  const kernel = [
      [1, 1, 1],
      [1, 1, 1],
      [1, 1, 1]
  ];
  const imageData = convert1Dto2D(image.data.map(pixel => pixel / 255), image.width);
  // 前向传播过程
  // 进行卷积操作
  const convOutput = convolve(imageData, kernel);
  // log('1.json', convOutput);
  // 进行最大池化操作
  const pooledOutput = maxPooling(convOutput);
  // log('2.json', pooledOutput);
  // 将池化后的二维数据扁平化
  const flattenedOutput = flatten(pooledOutput);
  // 通过全连接层得到最终的预测输出
  const fcOutput = fullyConnected(flattenedOutput, model.weights, model.biases);
  // 找到输出中概率最大的索引，即为预测的数字
  const prediction = fcOutput.indexOf(Math.max(...fcOutput));
  return prediction;
}

// 主函数
async function main() {
  const labels = await getUbyteFileData('./MINIST/t10k-labels-idx1-ubyte.gz');
  const images = await getUbyteFileData('./MINIST/t10k-images-idx3-ubyte.gz');

  // 保存模型
  const modelFilePath = 'model.json';

  // 加载模型
  const loadedModel = loadModel(modelFilePath);

  // 预测的图片数据
  // for (let i = 0; i < images.length; i++) {
  for (let i = 0; i < 1; i++) {
    const image = images[i];
    const prediction = predict(image, loadedModel);
    console.log(`预测的数字: ${prediction}, 实际是：${labels[i]}`);
  }
}
main();
